import json
import networkx as nx
import matplotlib.pyplot as plt
from networkx.drawing.nx_agraph import graphviz_layout

def draw_causal_graph(json_file_path, font_name='Loma'):
    # Read JSON file
    with open(json_file_path, 'r') as file:
        data = json.load(file)

    # Create a directed graph
    G = nx.DiGraph()

    # Add nodes and edges
    for source, targets in data.items():
        for target in targets:
            G.add_edge(source, target)  # Ensure correct node labels

    # Use Graphviz to create a hierarchical layout with increased spacing
    pos = graphviz_layout(G, prog='dot', args='-Granksep=2 -Gnodesep=1')  # Adjust ranksep and nodesep as needed

    # Graph settings
    plt.figure(figsize=(12, 8))  # Set figure size
    nx.draw(G, pos, with_labels=True, node_size=1000, node_color='lightblue', font_size=8, font_weight='bold', font_family=font_name, arrowstyle='-|>', arrowsize=20)
    plt.title('Causal Diagram')

    # Save and show the graph
    plt.savefig('causal_diagram.png', format='png', bbox_inches='tight')
    plt.show()

# Call the function
draw_causal_graph('experiments/meituan_micro/model-based/run-2/explain-causal-graph.json')
